﻿#include <xy/alg/interpolator.h>
#include <xy/solver/integrator.h>

namespace xy
{

namespace solver
{

using namespace xy::math;
using namespace xy::alg;

integrator::integrator(std::size_t series, double stepTime) : series(series), stepTime(stepTime), currentTime(0)
{
}

gtimer integrator::shift(funcxxt, gtimer U)
{
    // 时间平移
    currentTime += stepTime;
    return U + 1;
}

multi_step::multi_step(std::size_t series, double stepTime)
    : integrator(series, stepTime), alpha(series + 1), beta(series + 1)
{
    alpha.set_zero();
    beta.set_zero();
}

gtimer multi_step::shift(funcxxt f, gtimer U)
{
    // 获取隐函数
    auto F = [&](const vecxd &u) -> vecxd {
        auto ret = u * alpha[0] - f(u, currentTime + stepTime) * stepTime * beta[0];
        for (std::size_t i = 0; i < series; i++)
            ret += *(U - i) * alpha[i + 1] - f(*(U - i), currentTime - i * stepTime) * stepTime * beta[i + 1];
        return ret;
    };

    // 注意一定要先求解再赋值，如果 U 变化，则 F 也会变化，导致方程不对
    // 牛顿法求解隐式方程组，使用 U 作为初值
    *(U + 1) = nonlinear::solve(F, *U);

    return integrator::shift(f, U);
}

adams_bashforth::adams_bashforth(std::size_t series, double stepTime) : multi_step(series, stepTime)
{
    // 计算 f 系数
    std::vector<double> x(series);
    for (std::size_t i = 0; i < series; i++)
        x[i] = i;

    // 取 fi = 1 ，其它为 0 ，单独计算 fi 的系数
    for (std::size_t i = 0; i < series; i++)
    {
        std::vector<double> f(series, 0);
        f[i] = 1;

        interpolator_newton<double> interpolator(x, f);
        beta[series - i] = interpolator.interpolate().integrate(series - 1, series);
    }

    // u 系数
    alpha[0] = 1;
    alpha[1] = -1;
}

gtimer adams_bashforth::shift(funcxxt f, gtimer U)
{
    vecxd ret(U->size());
    ret.set_zero();
    for (std::size_t i = 0; i < series; i++)
        ret += *(U - i) * alpha[i + 1] - f(*(U - i), currentTime - i * stepTime) * stepTime * beta[i + 1];

    // beta[0] == 0，可显式计算
    *(U + 1) = ret / -alpha[0];

    // 时间平移
    currentTime += stepTime;

    return U + 1;
}

adams_moulton::adams_moulton(std::size_t series, double stepTime) : multi_step(series, stepTime)
{
    // 计算 f 系数
    std::vector<double> x(series + 1);
    for (std::size_t i = 0; i < series + 1; i++)
        x[i] = i;

    // 取 fi = 1 ，其它为 0 ，单独计算 fi 的系数
    for (std::size_t i = 0; i < series + 1; i++)
    {
        std::vector<double> f(series + 1, 0);
        f[i] = 1;

        interpolator_newton<double> interpolator(x, f);
        beta[series - i] = interpolator.interpolate().integrate(series - 1, series);
    }

    // u 系数
    alpha[0] = 1;
    alpha[1] = -1;
}

bdfs::bdfs(std::size_t series, double stepTime) : multi_step(series, stepTime)
{
    matxd A(series + 1, series + 1);
    vecxd b(series + 1);
    A.set_zero();
    b.set_zero();

    // FD 公式得到的线性系统
    b[1] = 1;
    for (std::size_t i = 0; i < series + 1; i++)
        A(0, i) = 1;

    for (std::size_t i = 1; i < series + 1; i++)
    {
        for (std::size_t j = 1; j < series + 1; j++)
        {
            auto sign = (i % 2 == 0) ? 1 : -1;
            A(i, j) = sign * std::pow(j, i);
        }
    }

    // 求解并归一化
    linear_lu solver(std::move(A));
    solver.perform();
    alpha = solver.solve(b.transposed()).transpose();
    beta[0] = b[1] / alpha[0];
    alpha /= alpha[0];
}

one_step::one_step(std::size_t series, double stepTime, std::size_t p)
    : integrator(series, stepTime), A(series, series), weights(series), nodes(series), p(p)
{
    A.set_zero();
    weights.set_zero();
    nodes.set_zero();
}

gtimer one_step::extrapolate(funcxxt f, gtimer U)
{
    // 外插法：U-1 是初值，U 是单步计算解
    time_grid UList(3, *U);
    gtimer iU = UList.begin();

    // 回退时间并改变步长
    currentTime -= stepTime;
    stepTime /= 2;

    // 两步计算
    iU = shift(f, iU);
    iU = shift(f, iU);

    // 恢复步长
    stepTime *= 2;

    // 使用更精确解
    *(U + 1) = *iU + (*iU - *(U + 1)) / (std::pow(2, p) - 1);

    return U + 1;
}

ERK::ERK(std::size_t series, double stepTime, std::size_t p) : one_step(series, stepTime, p)
{
}

gtimer ERK::shift(funcxxt f, gtimer U)
{
    std::size_t n = (*U).size(); // 初值向量
    std::size_t N = n * series;  // 变量规模

    // 方程组求根
    auto Y = preshift(f, U);

    // 计算下一步
    *(U + 1) = *U;
    for (std::size_t i = 0; i < series; i++)
        *(U + 1) += Y[i] * weights[i] * stepTime;

    return integrator::shift(f, U);
}

std::vector<vecxd> ERK::preshift(funcxxt f, gtimer U)
{
    std::size_t n = (*U).size(); // 初值向量
    std::size_t N = n * series;  // 变量规模

    std::vector<vecxd> Y;
    for (std::size_t i = 0; i < series; i++)
    {
        auto x = *U;
        for (std::size_t j = 0; j < i; j++)
            x += Y[j] * A(i, j) * stepTime;

        Y.emplace_back(f(x, currentTime + stepTime * nodes[i]));
    }
    return Y;
}

IRK::IRK(std::size_t series, double stepTime, std::size_t p) : one_step(series, stepTime, p)
{
}

gtimer IRK::shift(funcxxt f, gtimer U)
{
    std::size_t n = (*U).size(); // 初值向量
    std::size_t N = n * series;  // 变量规模

    // 方程组求根
    auto Y = preshift(f, U);

    // 计算下一步
    *(U + 1) = *U;
    for (std::size_t i = 0; i < series; i++)
        for (std::size_t j = 0; j < n; j++)
            (*(U + 1))[j] += stepTime * weights[i] * Y[i * n + j];

    return integrator::shift(f, U);
}

vecxd IRK::preshift(funcxxt f, gtimer U)
{
    std::size_t n = (*U).size(); // 初值向量
    std::size_t N = n * series;  // 变量规模

    // 获取多元隐函数 Y = (y1 y2 ... ys)
    auto F = [&](const vecxd &Y) -> vecxd {
        vecxd ret(N);

        // i ：第 i 个方程
        // k ：表示 yk
        // j ：yk 的第 j 个分量
        for (std::size_t i = 0; i < series; i++)
        {
            vecxd sy = *U;
            for (std::size_t j = 0; j < n; j++)
                for (std::size_t k = 0; k < series; k++)
                    sy[j] += stepTime * A(i, k) * Y[k * n + j];

            vecxd fi = f(sy, currentTime + stepTime * nodes[i]);
            for (std::size_t j = 0; j < n; j++)
                ret[i * n + j] = Y[i * n + j] - fi[j];
        }
        return ret;
    };

    // 用 f(u,t) 作为初始向量
    vecxd Y(N);
    auto Y0 = f(*U, currentTime);
    for (std::size_t i = 0; i < series; i++)
        for (std::size_t j = 0; j < n; j++)
            Y[i * n + j] = Y0[j];

    return nonlinear::solve(F, Y);
}

classical_RK4::classical_RK4(double k) : ERK(4, k, 4)
{
    // RK 节点和权
    nodes = {0, 0.5, 0.5, 1};
    weights = {1.0 / 6, 1.0 / 3, 1.0 / 3, 1.0 / 6};

    A(1, 0) = 0.5;
    A(2, 1) = 0.5;
    A(3, 2) = 1;
}

ESDIRK6::ESDIRK6(double k) : IRK(6, k, 4)
{
    // RK 节点和权
    nodes = {0, 0.5, 83.0 / 250, 31.0 / 50, 17.0 / 20, 1};
    weights = {82889.0 / 524892, 0, 15625.0 / 83664, 69875.0 / 102672, -2260.0 / 8211, 0.25};

    A(1, 0) = 0.25;
    A(1, 1) = 0.25;

    A(2, 0) = 8611.0 / 62500;
    A(2, 1) = -1743.0 / 31250;
    A(2, 2) = 0.25;

    A(3, 0) = 5012029.0 / 34652500;
    A(3, 1) = -654441.0 / 2922500;
    A(3, 2) = 174375.0 / 388108;
    A(3, 3) = 0.25;

    A(4, 0) = 15267082809.0 / 155376265600;
    A(4, 1) = -71443401.0 / 120774400;
    A(4, 2) = 730878875.0 / 902184768;
    A(4, 3) = 2285395.0 / 8070912;
    A(4, 4) = 0.25;

    for (std::size_t i = 0; i < series; i++)
        A(5, i) = weights[i];
}

gtimer ESDIRK6::shift(const matxd &B, func1x g, linear_ptlu &solver, gtimer U)
{
    // 存放 phi 值，初始化为 *U
    std::vector<vecxd> Phi(series, *U);

    // 对角线循环
    for (std::size_t i = 1; i < series; i++)
    {
        // b = U^n + k sum_j^(i-1) a_ij [A Phi[j] + g(t + cj * k)]
        for (std::size_t j = 0; j < i; j++)
            Phi[i] += (B * Phi[j] + g(currentTime + nodes[j] * stepTime)) * A(i, j) * stepTime;

        // 使用线性求解
        Phi[i] = solver.solve(Phi[i].transposed()).transpose();
    }

    // 加权求和
    *(U + 1) = *U;
    for (std::size_t i = 0; i < series; i++)
        *(U + 1) += (B * Phi[i] + g(currentTime + nodes[i] * stepTime)) * weights[i] * stepTime;

    // 注意传空指针
    return integrator::shift(nullptr, U);
}

gauss_legendre_RK::gauss_legendre_RK(std::size_t s, double k) : IRK(s, k, 2 * s)
{
    // 获得首一 Gauss-Legendre 多项式
    polyd GL = polyd::monic_legendre(s);

    // 依次计算 s 个根，从 0 开始迭代，通常会得到从小到大的根
    for (std::size_t i = 0; i < s; i++)
    {
        // 构造函数
        func11 F = [&](double x) -> double { return GL.value(x); };

        // 迭代求根
        nodes[i] = nonlinear::solve(F, 0);

        // 除以因式 x - c[i] 进行降阶
        GL = GL / polyd({-nodes[i], 1});
    }

    // 保险起见按照升序排列
    std::sort(nodes.flatten().begin(), nodes.flatten().end());

    // 初始化权
    for (std::size_t i = 0; i < s; i++)
        weights[i] = polyd::lagrange(nodes.flatten(), i).integrate(0, 1);

    // 初始化系数矩阵
    for (std::size_t i = 0; i < s; i++)
        for (std::size_t j = 0; j < s; j++)
            A(i, j) = polyd::lagrange(nodes.flatten(), j).integrate(0, nodes[i]);
}

gtimer gauss_legendre_RK::shift(funcxxt f, gtimer U)
{
    std::size_t N = U->size();

    // 变量规模：共 series 个方程，每个方程是一个 N 元向量
    // Y = (y1 y2 ... ys)
    funcxx F = [&](const vecxd &Y) -> vecxd {
        // 结果容器
        vecxd ret(Y.size());

        // i ：第 i 个方程
        for (std::size_t i = 0; i < series; i++)
        {
            // 获得初始向量
            vecxd yi = *U;

            for (std::size_t j = 0; j < series; j++)
                for (std::size_t n = 0; n < N; n++)
                    yi[n] += stepTime * A(i, j) * Y[j * N + n];

            // 计算 yi
            yi = f(yi, currentTime + stepTime * nodes[i]);

            // 得到第 i 方程
            for (std::size_t n = 0; n < N; n++)
                ret[i * N + n] = Y[i * N + n] - yi[n];
        }
        return ret;
    };

    // 使用 f(U^n,t) 作为初值
    vecxd Y0 = f(*U, currentTime);
    vecxd Y(series * N);

    for (std::size_t i = 0; i < series; i++)
        for (std::size_t n = 0; n < N; n++)
            Y[i * N + n] = Y0[n];

    // 牛顿法迭代求根
    Y = nonlinear::solve(F, Y);

    // 加权推进
    *(U + 1) = *U;
    for (std::size_t i = 0; i < series; i++)
        for (std::size_t n = 0; n < N; n++)
            (*(U + 1))[n] += stepTime * weights[i] * Y[i * N + n];

    return integrator::shift(f, U);
}

embedded_RK::embedded_RK(std::size_t series, double stepTime, std::size_t p, std::size_t dp)
    : ERK(series, stepTime, p), dualWeights(series), dp(dp)
{
    dualWeights.set_zero();
}

gtimer embedded_RK::auto_step(funcxxt f, gtimer U)
{
    std::size_t n = (*U).size(); // 初值向量规模
    std::size_t N = n * series;  // 变量规模
    vecxd Un(n);

    constexpr double Eabs = 1e-5;    // 绝对误差
    constexpr double Erel = 1e-5;    // 相对误差
    constexpr double rho = 0.8;      // rho = 0.8 0.9
    constexpr double rhomax = 1.5;   // rhomax [1.5,5]
    constexpr double rhomin = 0.2;   // rhomin [0.2,rho]
    std::size_t q = std::min(p, dp); // 确定放大指数

    double t = currentTime; // 记录时间
    double Eind = 2;        // 误差指示器
    while (Eind > 1)
    {
        // 方程组求根
        auto Y = preshift(f, U);

        // 计算两个解
        auto dualUn = Un = *U;
        for (std::size_t i = 0; i < series; i++)
        {
            Un += Y[i] * weights[i] * stepTime;
            dualUn += Y[i] * dualWeights[i] * stepTime;
        }

        // 记录移动时间
        t = currentTime + stepTime;

        // 计算误差 Eind
        Eind = 0;
        for (std::size_t i = 0; i < n; i++)
            Eind += std::pow((Un[i] - dualUn[i]) / (Eabs + std::abs((*U)[i]) * Erel), 2);
        Eind = std::sqrt(Eind / n);

        // 修正步长
        stepTime *= std::min(rhomax, std::max(rhomin, rho * std::pow(Eind, -1.0 / (q + 1))));
    }

    *(U + 1) = Un;   // 给出解
    currentTime = t; // 移动时间

    return U + 1;
}

fehlberg6::fehlberg6(double k) : embedded_RK(6, k, 4, 5)
{
    // RK 节点和权
    nodes = {0, 0.25, 3.0 / 8, 12.0 / 13, 1, 0.5};
    weights = {25.0 / 216, 0, 1408.0 / 2565, 2197.0 / 4104, -0.2, 0};
    dualWeights = {16.0 / 135, 0, 6656.0 / 12825, 28561.0 / 56430, -9.0 / 50, 2.0 / 55};

    A(1, 0) = 0.25;

    A(2, 0) = 3.0 / 32;
    A(2, 1) = 9.0 / 32;

    A(3, 0) = 1932.0 / 2197;
    A(3, 1) = -7200.0 / 2197;
    A(3, 2) = 7296.0 / 2197;

    A(4, 0) = 439.0 / 216;
    A(4, 1) = -8;
    A(4, 2) = 3680.0 / 513;
    A(4, 3) = -845.0 / 4104;

    A(5, 0) = -8.0 / 27;
    A(5, 1) = 2;
    A(5, 2) = -3544.0 / 2565;
    A(5, 3) = 1859.0 / 4104;
    A(5, 4) = -11.0 / 40;
}

dormand_prince7::dormand_prince7(double k) : embedded_RK(7, k, 5, 4)
{
    // RK 节点和权
    nodes = {0, 0.2, 0.3, 0.8, 8.0 / 9, 1, 1};
    weights = {35.0 / 384, 0, 500.0 / 1113, 125.0 / 192, -2187.0 / 6784, 11.0 / 84, 0};
    dualWeights = {5179.0 / 57600, 0, 7571.0 / 16695, 393.0 / 640, -92097.0 / 339200, 187.0 / 2100, 1.0 / 40};

    A(1, 0) = 0.2;

    A(2, 0) = 3.0 / 40;
    A(2, 1) = 9.0 / 40;

    A(3, 0) = 44.0 / 45;
    A(3, 1) = -56.0 / 15;
    A(3, 2) = 32.0 / 9;

    A(4, 0) = 19372.0 / 6561;
    A(4, 1) = -25360.0 / 2187;
    A(4, 2) = 64448.0 / 6561;
    A(4, 3) = -212.0 / 729;

    A(5, 0) = 9017.0 / 3168;
    A(5, 1) = -355.0 / 33;
    A(5, 2) = 46732.0 / 5247;
    A(5, 3) = 49.0 / 176;
    A(5, 4) = -5103.0 / 18656;

    for (std::size_t i = 0; i < series; i++)
        A(6, i) = weights[i];
}

} // namespace solver

} // namespace xy